import pandas as pd
import numpy as np
import math
import glob
from tqdm import tqdm





#height = [150,200,250,300,350,400]
height = 200
pointcloud_dir = "H:/osdata/STPLS3D/Synthetic_v3_InstanceSegmentation"
out_pos_file_dir = "H:/osdata/STPLS3D/ws/Synthetic_v3_campos_val/" + str(height) + "/"

scene_id = 1

def get_instance_centerpoints(allpoints):
    sem_ins_labels = np.ascontiguousarray(allpoints[:, -2:])
    sem_ins_labels = sem_ins_labels.astype(np.int32)
    
    building_instance_labels = np.where(sem_ins_labels[:,0] == 1,sem_ins_labels[:,1],-100)
    building_instance_labels = np.unique(building_instance_labels)[1:]

    # building_instance_points_index = np.where(sem_ins_labels[:,0] == 1)

    ins_centre_xyz = {}
    for id in building_instance_labels:
        index = np.where(sem_ins_labels[:,1] == id)
        points = np.ascontiguousarray(allpoints[index])
        xyz = np.mean(points[:,:3],axis=0)
        ins_centre_xyz[id] = xyz


    return ins_centre_xyz

def get_cam_center_points(allpoints):
    all_xy = np.ascontiguousarray(allpoints[:, 0:2])
    # sem_ins_labels = all_xy.astype(np.int32)
    
    minx = np.min(all_xy[:,0])+10
    miny = np.min(all_xy[:,1])+10

    maxx = np.max(all_xy[:,0])-10
    maxy = np.max(all_xy[:,1])-10

    points = []
    for x in range(5):
        nx = x / 5.0
        for y in range(5):
            ny = y / 5.0
            px = minx * (1-nx) + maxx * nx
            py = miny * (1-ny) + maxy * ny
            points.append([px,py])
    return points

def get_cam_pos(center,height,heading,pitch):
    
    heading = heading/180 * math.pi
    pitch = pitch/180 * math.pi
    abs_h = height - center[2]
    xy_r = math.tan(math.pi/2 + pitch) * abs_h * -1
    cam_y = xy_r * math.cos(heading)
    cam_x = xy_r * math.sin(heading)

    cam_x += center[0]
    cam_y += center[1]

    return cam_x,cam_y


def gen_pos(points):

    heading = [0,90,180,270]
    pitch = [-45]

    res_cam_pose = []
    for xy in points:
        for p in pitch:
            for h in heading:
                xyz = [xy[0],xy[1],0]
                cx,cy = get_cam_pos(xyz,height,h,p)
                res_cam_pose.append([cx,cy,height,h,p,0,scene_id])
    return res_cam_pose

def write_file(out_pos_file_path,res_cam_pose):

    out_pos_file = open(out_pos_file_path,'w+')
    out_pos_file.write("#x y z heading pitch roll scene_id\n")
    for pose in res_cam_pose: 
        str2 = '@'.join([str(f) for f in pose])
        out_pos_file.write(str2+"\n")


def gen_pose_file(path,outpath):
    points = pd.read_csv(path, header=None).values
    centers = get_cam_center_points(points)
    cam_pose = gen_pos(centers)
    write_file(outpath,cam_pose)




files = glob.glob(pointcloud_dir + "/*.txt")
pbar = tqdm(total=len(files))
for file in files:
    basefilename = file.replace(pointcloud_dir,"")
    scene_id = basefilename.replace("_points_GTv3.txt","")
    scene_id = scene_id.replace("\\","")
    scene_id = int(scene_id)
    outp = out_pos_file_dir + basefilename
    gen_pose_file(file,outp)
    pbar.update()



